
from struct import *
from primitiveBase import *

class genBase(object):
    '''
    base class for generators
    '''
    MINSTRINGSIZE =  3 
    SIZER_VARIANCE = 2
    fileId = 0
    #search for sizers after a strings location not just before the location
    SIZER_SEARCH_AFTER = True 
    scriptLegos = []
    
    def __init__(self):
        '''
        Constructor
        '''
        pass

    def writeToFile(self, scriptCode, outputDir):
        fname = "%s/%s%d.py" % (outputDir, self.__class__.__name__, self.fileId)
        tmpFile = open(fname, "w+b")
        tmpFile.write(scriptCode);
        tmpFile.close()
        self.fileId += 1
        return fname

    def writeFuzzScripts(self, inputBuff=None):
        pass
    """
    Given a value, variance and data will scan data for values within the range
    Returns the location of a value, the value, and the perceived type
    """
    def findValues(self, value, variance, inputBuff, types = (), name=""):
        
        w16 = word16()
        w32 = word32()
        w64 = word64()
        wordTypes = (w64, w32, w16)
        
        values = range(value - variance, value + variance +1)
            
        foundValues = []
        foundValuesLocations = set()

        for wordType in wordTypes:
            for x in range(0, len(inputBuff) - wordType.typeSize + 1):
                #unpack and compare to the value
                #TODO test to see if its faster to pack then search, rather than unpack.
                possibleLocation = set(range(x, x+wordType.typeSize))
                if possibleLocation <= foundValuesLocations:
                    break
                for unpackFormat in wordType.unpackFormats:
                    unpackString = inputBuff[x:x+calcsize(unpackFormat)]
                    format = "<" + unpackFormat #try little endian
                    result = unpack(format, unpackString)[0]
                    #print result
                    if result in values:
                        variance = result - value
                        foundValue = wordType.__new__(type(wordType))
                        foundValue.__init__(endian="<", originalValue=result, streamOffset=x, name=name, variance=variance)
                        foundValues.append(foundValue)
                        foundValuesLocations.update(possibleLocation)
                        break
                        
                    format = ">" + unpackFormat #try bigendian
                    result = unpack(format, unpackString)[0]
                    if result in values:
                        variance = result - value
                        foundValue = wordType.__new__(type(wordType))
                        foundValue.__init__(endian="<", originalValue=result, streamOffset=x, name=name, variance=variance)
                        foundValues.append(foundValue)
                        foundValuesLocations.update(possibleLocation)
                        break
        return foundValues
    
    """
    charSetContainsNumeric - when true the SIZER_VARIANCE is increased by the magnitude of the length, because the string found could contain the sizer itself
    sizeInBlock - allows the sizer to be within the block
    """
    def findSizers(self, inputBuff, start, end, blockName, sizerInBlock=True, charSetContainsNumeric=False):
        #given a list of ranges this will find potential sizers
        #create a cache of values and their locations
        if charSetContainsNumeric:
            variance = self.SIZER_VARIANCE + int((end - start) / 10)
        else:
            variance = self.SIZER_VARIANCE
        words = self.findValues(end-start, variance, inputBuff)
        #convert each value into a sizer and add the block name 
        sizers = []
        for word in words:
            s = sizer(word)
            s.setBlockName(blockName)
            sizers.append(s)
        return sizers
    
    
    def findStringsInBinary(self, inputBuff):
        if len(inputBuff) < self.MINSTRINGSIZE:
            return []
        stringsFound = []
        stringsFoundLoc = []

        for i in range(0, len(inputBuff)):
            #for this loop we dont want to re-analyze sections after they are already in the stringsFound list
            #other search loops may want to re-analyze sections already in the results list
            skipIndex = False
            for (start, end) in stringsFoundLoc:
                if start <= i and end >= i:
                    skipIndex = True
                    break
            if skipIndex:
                continue
            
            if inputBuff[i] in string.ascii_letters and inputBuff[i+self.MINSTRINGSIZE] in string.ascii_letters:
                stringFound = True
                for x in range(i+1, i+self.MINSTRINGSIZE-1):
                    if inputBuff[x] not in string.ascii_letters:
                        stringFound = False
                if stringFound:
                    stringStart = i
                    stringEnd = -1
                    #find the end of the string, possibly adding substrings to the list based on delims other than non-ascii chars
                    for x in range(i, len(inputBuff)):
                        if inputBuff[x] not in string.ascii_letters:
                            stringEnd = x
                            break
                    if x == len(inputBuff)-1:
                        stringEnd = x+1
                    if stringEnd == -1:
                        raise Exception('stringEnd not found in findStringsInBinary')
                    m = md5()
                    m.update(inputBuff[stringStart:stringEnd])
                    name = m.hexdigest()
                    foundString = stringPrimitive(str(inputBuff[stringStart:stringEnd]), stringStart, name=name)
                    stringsFound.append(foundString)
                    stringsFoundLoc.append((stringStart, stringEnd))
                    
        return stringsFound
